import matplotlib.pyplot as plt

# Define a function to read data
def read_data(file_path):
    x_coords = []
    y_coords = []
    with open(file_path, 'r') as file:
        for line in file:
            x, y = map(float, line.split())
            x_coords.append(x)
            y_coords.append(y)
    return x_coords, y_coords

# Define a function to plot scatter points and lines
def plot_data(x_coords, y_coords, label, color):
    plt.scatter(x_coords, y_coords, color=color, label=label, s=5)  # Reduce the size of the points
    plt.plot(x_coords, y_coords, color=color)  # Add lines

# File paths
file_paths = {
    'B': ['output/B_2.txt', 'output/B_4.txt', 'output/B_6.txt', 'output/B_8.txt'],
    'C': ['output/C_5.txt', 'output/C_10.txt', 'output/C_15.txt', 'output/C_20.txt'],
    'E': ['output/E_1.txt', 'output/E_2.txt'],
    'F': ['output/F_10.txt', 'output/F_40.txt', 'output/F_160.txt']
}

# Define colors and labels
colors = ['r', 'g', 'b', 'y']

# Plot files starting with B
for key in ['B']:
    labels = ['n=2', 'n=4', 'n=6', 'n=8']
    fig, ax = plt.subplots()
    for i, file_path in enumerate(file_paths[key]):
        x_coords, y_coords = read_data(file_path)
        plot_data(x_coords, y_coords, labels[i], colors[i])
    ax.legend()
    ax.set_title(r'Newton Interpolation on $\frac{1}{1+x^2}$, Equally Spaced Points')
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(f'images/{key}_plot.png')
    
# Plot files starting with C
for key in ['C']:
    labels = ['n=5', 'n=10', 'n=15', 'n=20']
    fig, ax = plt.subplots()
    for i, file_path in enumerate(file_paths[key]):
        x_coords, y_coords = read_data(file_path)
        plot_data(x_coords, y_coords, labels[i], colors[i])
    ax.legend()
    ax.set_title(r'Newton Interpolation on $\frac{1}{1+25x^2}$, Chebyshev Points')
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(f'images/{key}_plot.png')
    
# Plot files starting with E
for key in ['E']:
    labels = ['sp1', 'sp2']
    fig, ax = plt.subplots()
    for i, file_path in enumerate(file_paths[key]):
        x_coords, y_coords = read_data(file_path)
        # Only select the first 31 lines of data for plotting
        x_coords, y_coords = x_coords[:310], y_coords[:310]
        plot_data(x_coords, y_coords, labels[i], colors[i])
    ax.legend()
    ax.set_title(r'Weight change of two samples')
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(f'images/{key}_plot.png')

# Plot files starting with F
for i, file_path in enumerate(file_paths['F']):
    fig, ax = plt.subplots()
    x_coords, y_coords = read_data(file_path)
    plot_data(x_coords, y_coords, f'F_{i+1}', colors[i])
    ax.legend()
    ax.set_title(f'Scatter Plot from {file_path}')
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    # Use part of the file name to generate the saved file name
    key = file_path.split('_')[1].split('.')[0]
    plt.savefig(f'images/F_{key}_plot.png')